#!/usr/bin/env python
'''
Given the TranSEQ output and the contigs.fa
This program select the contigs that 
(1) The longest one
(2) The contig with abundance level over some threshold
(3) Those with Inf likelihood ratio
'''
from parse_fasta_file import get_seqs_len
import string
import os
import getopt
import sys

def get_longest_seq_in_fa( faFileN ):
    ''' Given a fasta, get the longest seq's name '''
    max_seqL = 0
    longest_seqN = ""
    seqsL = get_seqs_len( faFileN )
    for seqN, L in seqsL.items():
        if( max_seqL < L ):
            max_seqL = L
            longest_seqN = seqN
    return longest_seqN

def get_surly_expressed_seq_in_fa( TranSEQoutFN ):
    ''' Given TranSEQ output, find the seqNs
    with Inf log LR for expression.
    '''
    seqNs = []
    NO_COL_IN_TranSEQ_OUT = 3
    ifile = open ( TranSEQoutFN, "r" )
    for line in ifile.readlines():
        line = line.strip()
        record = line.split( '\t' )
        if( len( record ) == NO_COL_IN_TranSEQ_OUT ):
            if record[2] == 'Inf':
                seqNs.append( record[0] )
    ifile.close()
    return seqNs

def get_unlikely_expressed_seq_in_fa( TranSEQoutFN ):
    ''' Given TranSEQ output, find the seqNs
    with negative log LR for expression.
    '''
    seqNs = []
    NO_COL_IN_TranSEQ_OUT = 3
    ifile = open ( TranSEQoutFN, "r" )
    for line in ifile.readlines():
        line = line.strip()
        record = line.split( '\t' )
        if( len( record ) == NO_COL_IN_TranSEQ_OUT ):
            exp = record[1].strip()
            logLR = record[2].strip()
            if ( logLR.startswith( '-' ) \
                or ( exp == '0' )
                or ( ( 'e' in exp ) and ( 'e' in logLR  or logLR == 0 ) ) ):
                    seqNs.append( record[0] )
    ifile.close()
    return seqNs

def format_num( numStr ):
    '''Return the string better format the scientific num'''
    precision2digit = 3
    if( numStr == 'Inf' ):
        return( numStr )
    elif( 'e' in numStr[2:] ):
        return( '0' )
        #return( numStr[0] + numStr[numStr.find( "e" ):] )
    elif( numStr.startswith( '0.' ) and len( numStr ) > precision2digit + 1 ):
        return( numStr[1:( precision2digit + 2 )] )
    elif( not numStr.startswith( '0' ) and ( '.' in numStr ) ):
        return( numStr[:numStr.find( "." )] ) #return integer part
    else:
        return( numStr )

def get_seqs_info( TranSEQOutFN ):
    ''' translate TranSEQ output to dict structure '''
    d = dict()
    NO_COL_IN_TranSEQ_OUT = 3
    ifile = open ( TranSEQOutFN, "r" )
    line = ifile.readline() # remove the first line
    for line in ifile.readlines():
        line = string.strip( line )
        record = line.split( '\t' )
        if len( record ) == NO_COL_IN_TranSEQ_OUT:
            d[record[0]] = format_num( record[1] )\
                + '_' + format_num( record[2] )
    ifile.close()
    return d

def get_highly_expressed_seq_in_fa( TranSEQoutFN, threshold ):
    ''' Given TranSEQ output, find the seqs 
     which have expression levels over the threshold.
    '''
    seqs = []
    NO_COL_IN_TranSEQ_OUT = 3
    ifile = open ( TranSEQoutFN, "r" )
    for line in ifile.readlines():
        line = line.strip()
        record = line.split( '\t' )
        if( len( record ) == NO_COL_IN_TranSEQ_OUT ):
            if( record[1].startswith( '0.' ) and \
                float( record[1] ) >= threshold ) or \
                ( record[1] == '1' ):
                seqs.append( record[0] )
    ifile.close()
    return seqs

def print_selected_seqs( selected_seqsN, faFileN, outFileN, seqs_info_dic ):
    ifile = open( faFileN, 'r' )
    ofile = open( outFileN, 'w' )
    flag = False # flag to indicate if a sequence is selected to out
    while ( 1 ):
        line = ifile.readline()
        if not line:
            break;
        elif ( line[0] == '>' ):
            seqN = line.split()[0][1:]
            if( seqN in selected_seqsN ):
                flag = True
                if seqs_info_dic.has_key( seqN ):
                    ofile.write( '>' + seqN + '_' + seqs_info_dic[seqN] + '\n' )
                else :
                    ofile.write( '>' + seqN + '\n' )
            else:
                flag = False
        else:
            if flag:
                ofile.write( line )
    ifile.close()
    ofile.close()

def select_contigs( faFileN, TranSEQoutFN, outFileN ):
    ''' Print the selected contigs from the result of TranSEQ
    to the output file.'''
    EXP_THRESHOLD = 0.001
    selected_seqsN = set();
    selected_seqsN.add( get_longest_seq_in_fa( faFileN ) )
    selected_seqsN.update( \
        get_surly_expressed_seq_in_fa( TranSEQoutFN ) )
    selected_seqsN.update( \
        get_highly_expressed_seq_in_fa( TranSEQoutFN, EXP_THRESHOLD ) )
    selected_seqsN = selected_seqsN.difference\
        ( get_unlikely_expressed_seq_in_fa( TranSEQoutFN ) )
    seqs_dic = get_seqs_info( TranSEQoutFN )
    print_selected_seqs( selected_seqsN, faFileN, outFileN, seqs_dic )
    print str( len( selected_seqsN ) ) + "CDSs are selected"
    return( selected_seqsN )

def usage():
    sys.stderr.write( "Working dir is " + os.getcwd() + '\n' )
    print __doc__
    print "usage: ", sys.argv[0]
    print "The three parameters should be "
    print "(1) contigs file"
    print "(2) Result from TranSEQ"
    print "(3) output file"
    sys.exit( 2 )

def main():
    try:
        opts, args = getopt.getopt( sys.argv[1:], 'h', ["help"] )
    except getopt.error:
        usage()
        sys.exit( 2 );
    # process options and argument 
    for o in opts:
        if o in ( "-h", "--help" ):
            usage()
            sys.exit();
    argc = len( args )
    if( argc == 3 ):
        select_contigs( args[0], args[1], args[2] )
    else:
        usage()
        print str( argc ) + " args are given"

if __name__ == '__main__':
    main()
